import argparse
import logging
import math
import os

import numpy
import numpy as np
from deap import tools

from plotting.drawer import draw_precedence_relations, plot_gantt_chart
from solution_methods.HHO.src.initialization import initialize_run
from solution_methods.HHO.src.operators import evaluate_individual, variation, repair_precedence_constraints, \
    evaluate_population
from solution_methods.HHO.utils import record_stats, output_dir_exp_name, results_saving
from solution_methods.helper_functions import load_parameters, load_job_shop_env

PARAM_FILE = "../../configs/HHO.toml"


def run_HHO(jobShopEnv, population, toolbox, stats, hof, **kwargs):
    """
    执行Harris Hawks Optimization算法并返回最佳个体。

    参数:
        jobShopEnv: 问题环境。作业车间问题环境
        population: 初始种群。
        toolbox: DEAP工具箱。
        stats: DEAP统计信息。
        hof: Hall of Fame（名人堂）。
        kwargs: 额外的关键字参数。

    返回:
        tuple: 最佳个体的makespan和作业车间环境。
    """

    # 初始化名人堂和统计信息
    hof.update(population)

    gen = 0
    logbook = tools.Logbook()
    logbook.header = ["gen"] + (stats.fields if stats else [])
    # 如果 stats 存在且有 fields 属性，则返回 stats.fields；否则返回空列表 []。
    df_list = []

    # initialize the location and Energy of the rabbit
    # 初始化兔子位置 和 兔子能量
    # dim = 30  # 假设问题维度为30
    # dim = kwargs['algorithm'].get('dim')
    pop_size = kwargs['algorithm'].get('population_size')
    # ub = kwargs['algorithm'].get('ub')
    # lb = kwargs['algorithm'].get('lb')

    # lb 和 ub现在是机器MS码段的概率系数 的边界，所以解空间是从0-1，后面乘以 每个操作的可选机器数  然后  按照四舍五入保留整数 就行
    lb = 0
    ub = 1

    dim = len(jobShopEnv.operations_machine_num) # 问题规模dim，等于所有工序数

    # # 遍历population数组的第一个list，求出第一个list的平均位置
    # AVG_Location = numpy.mean(population[0], axis=0)

    # Rabbit_Location = numpy.zeros(dim)
    Rabbit_Location = [] # 当前最佳适应度个体

    Rabbit_Energy = float("inf")  # change this to -inf for maximization problems

    # 记录初始种群的统计信息
    record_stats(gen, population, logbook, stats, kwargs['output']['logbook'], df_list, logging)
    if kwargs['output']['logbook']:  # 判断配置中是否开启了日志记录功能
        logging.info(logbook.stream)

    for gen in range(1, kwargs['algorithm']['ngen'] + 1):

        ngen = kwargs['algorithm']['ngen']  # 从配置文件中读取到的最大迭代次数

        for i in range(1, ngen + 1):

            # Check boundries

            # X[i, :] = numpy.clip(X[i, :], lb, ub)


            # 解码
            # for int

            # fitness of locations
            # 调用目标函数，返回适应度值,评估种群的适应度
            fitness = evaluate_population(toolbox, population)

            # 如果个体适应度值小于兔子能量，则更新兔子位置和能量，兔子就是当前的最优解个体
            for i in range(len(population)):
                if fitness[i][0] < Rabbit_Energy:  # Change this to > for maximization problem
                    Rabbit_Energy = fitness[i][0]
                    Rabbit_Location = population[i]

            b = Rabbit_Location

        E1 = 2 * (1 - (gen / ngen))  # factor to show the decreaing energy of rabbit 兔子逃逸能量递减的系数

        # 种群针对HHO编码
        population = encode(jobShopEnv.operations_machine_num, population)

        # 哈里斯鹰捕获兔子，更新种群
        offspring = variation(population, toolbox, E1, pop_size,
                              ub, lb, dim, Rabbit_Location)

        # 种群针对HHO解码
        offspring = decode(jobShopEnv.operations_machine_num, offspring)

        a = offspring
        a = a


        # 如果环境需要，修复优先约束关系（仅适用于装配调度）
        if any(keyword in jobShopEnv.instance_name for keyword in ['/dafjs/', '/yfjs/']):
            try:
                offspring = repair_precedence_constraints(jobShopEnv, offspring)
            except Exception as e:
                logging.error(f"修复优先约束关系时出错: {e}")
                continue


        # 评估后代的适应度
        try:
            # 种群针对HHO解码
            # offspring = decode(jobShopEnv.operations_machine_num, offspring)
            a = offspring
            fitnesses = evaluate_population(toolbox, offspring)
            for ind, fit in zip(offspring, fitnesses):
                ind.fitness.values = fit
        except Exception as e:
            logging.error(f"评估后代适应度时出错: {e}")
            continue

        # 选择下一代种群
        population[:] = toolbox.select(population + offspring)

        # 更新名人堂和统计信息
        hof.update(population)
        record_stats(gen, population, logbook, stats, kwargs['output']['logbook'], df_list, logging)

    # 评估最佳个体的makespan
    makespan, jobShopEnv = evaluate_individual(hof[0], jobShopEnv, reset=False)
    logging.info(f"Makespan: {makespan}")
    return makespan, jobShopEnv

    #     # 更新名人堂和统计信息
    #     hof.update([creator.Individual(X[i, :]) for i in range(SearchAgents_no)])
    #     record_stats(t+1, [creator.Individual(X[i, :]) for i in range(SearchAgents_no)], logbook, stats, kwargs['output']['logbook'], df_list, logging)
    #
    #     t += 1
    #
    # timerEnd = time()
    # s['endTime'] = strftime("%Y-%m-%d-%H-%M-%S")
    # s['executionTime'] = timerEnd - timerStart
    # s['best'] = Rabbit_Energy
    # s['bestIndividual'] = Rabbit_Location
    #
    # # 评估最佳个体的makespan
    # makespan, jobShopEnv = evaluate_individual(hof[0], jobShopEnv, reset=False)
    # logging.info(f"Makespan: {makespan}")
    # return makespan, jobShopEnv

# MS机器选择部分，编码
def encode(operations_machine_num, population):
    for i in range(len(population)):
        ind_temp = []
        for j in range(len(operations_machine_num)):
            ind_temp.append(population[i][0][j] / operations_machine_num[j])
        population[i][0] = ind_temp
    return population

# MS机器选择部分，解码
def decode(operations_machine_num, population):
    for i in range(len(population)):
        ind_temp = []
        for j in range(len(operations_machine_num)):
            ind_temp.append(round(population[i][0][j] * operations_machine_num[j]))
        population[i][0] = ind_temp
    return population

def Levy(dim):
    beta = 1.5
    sigma = (
                    math.gamma(1 + beta)
                    * math.sin(math.pi * beta / 2)
                    / (math.gamma((1 + beta) / 2) * beta * 2 ** ((beta - 1) / 2))
            ) ** (1 / beta)
    u = 0.01 * np.random.randn(dim) * sigma
    v = np.random.randn(dim)
    zz = np.power(np.absolute(v), (1 / beta))
    step = np.divide(u, zz)
    return step


# # 辅助函数：用于记录每一代的统计信息
# def record_stats(gen, population, logbook, stats, output_logbook, df_list, logging):
#     if stats:
#         record = stats.compile(population) if stats else {}
#         logbook.record(gen=gen, **record)
#         if output_logbook:
#             logging.info(logbook.stream)


def main(param_file=PARAM_FILE):
    """
    主函数，加载参数文件，初始化作业车间环境和HHO算法，并运行HHO算法。

    参数:
        param_file: 参数文件路径，默认为PARAM_FILE。
    """
    try:
        parameters = load_parameters(param_file)
        logging.info(f"参数已从 {param_file} 加载。")
    except FileNotFoundError:
        logging.error(f"未找到参数文件 {param_file}。")
        return

    # 加载作业车间环境，并初始化HHO哈里斯鹰算法
    jobShopEnv = load_job_shop_env(parameters['instance'].get('problem_instance'))
    population, toolbox, stats, hof = initialize_run(jobShopEnv, **parameters)
    makespan, jobShopEnv = run_HHO(jobShopEnv, population, toolbox, stats, hof, **parameters)

    if makespan is not None:
        # 检查输出配置并准备输出路径（如果需要）
        output_config = parameters['output']
        save_gantt = output_config.get('save_gantt')
        save_results = output_config.get('save_results')
        show_gantt = output_config.get('show_gantt')
        show_precedences = output_config.get('show_precedences')

        if save_gantt or save_results:
            output_dir, exp_name = output_dir_exp_name(parameters)
            output_dir = os.path.join(output_dir, f"{exp_name}")
            os.makedirs(output_dir, exist_ok=True)

        # 绘制优先关系图（如果需要）
        if show_precedences:
            draw_precedence_relations(jobShopEnv)

        # 绘制甘特图（如果需要）
        if show_gantt or save_gantt:
            logging.info("正在生成甘特图。")
            plt = plot_gantt_chart(jobShopEnv)

            if save_gantt:
                plt.savefig(output_dir + "/gantt.png")
                logging.info(f"甘特图已保存到 {output_dir}")

            if show_gantt:
                plt.show()

        # 保存结果（如果启用）
        if save_results:
            results_saving(makespan, output_dir, parameters)
            logging.info(f"结果已保存到 {output_dir}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="运行HHO算法")
    parser.add_argument(
        "-f", "--config_file",
        type=str,
        default=PARAM_FILE,
        help="配置文件路径",
    )
    args = parser.parse_args()
    main(param_file=args.config_file)
